#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from pymongo import MongoClient
import os
import sys


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def GenerateFasta(collection, record):
    id = record.id.split(".")[0]
    version = record.id.split(".")[1]
    description = record.description[record.description.find(" ") + 1:]
    fasta = record.seq

    return_records = []

    for result in collection.find(
            {"gene_commentary_accession": id, "gene_commentary_version": version, "seq_id_gi": {"$exists": 1}}):
        if "seq_interval_from" in result:
            strand = "plus"
            if "na_strand" in result:
                if result["na_strand"] == "minus":
                    strand = "minus"
            gi = result["gene_commentary_accession"]
            new_id = ""
            new_description = ""
            new_seq = ""

            description_add = "|"
            if "gene_ref_locus" in result:
                description_add += result["gene_ref_locus"]
            if "prot_ref_name_E" in result:
                if description_add == "|":
                    description_add += result["prot_ref_name_E"]
                else:
                    description_add += " " + result["prot_ref_name_E"]
            if "org_ref_taxname" in result:
                if description_add == "|":
                    description_add += "[" + result["org_ref_taxname"] + "]"
                else:
                    description_add += " [" + result["org_ref_taxname"] + "]"
            if description_add != "|":
                new_description = description + description_add
            else:
                new_description = description

            if strand == "plus":
                new_id = id + "." + version + ":" + str(int(result["seq_interval_from"]) + 1) + "-" + str(
                    int(result["seq_interval_to"]) + 1) + "|EntrezGeneID:" + result["gene_track_geneid"]
                new_seq = fasta[int(result["seq_interval_from"]):int(result["seq_interval_to"]) + 1]
                new_record = SeqRecord(new_seq, id=new_id, description=new_description)
                if len(new_record.seq) > 0:
                    return_records.append(new_record)
                else:
                    print("No fasta seq " + new_id + " " + new_description)
            else:
                new_id = id + "." + version + ":c" + str(int(result["seq_interval_to"]) + 1) + "-" + str(
                    int(result["seq_interval_from"]) + 1) + "|EntrezGeneID:" + result["gene_track_geneid"]
                new_record = record
                new_record.seq = fasta[int(result["seq_interval_from"]):int(result["seq_interval_to"]) + 1]
                new_record = new_record.reverse_complement(id=new_id, description=new_description)
                if len(new_record.seq) > 0:
                    return_records.append(new_record)
                else:
                    print("No fasta seq " + new_id + " " + new_description)
    return return_records


def GenerateFile(inputfilepath):
    client, collection = InitMongoDB("localhost", 27017, "Pre_Data", "MicrobiomeGene")

    abspath = os.path.abspath(inputfilepath)
    parentpath = abspath[:abspath.rindex("/")]
    basename = os.path.basename(abspath)
    outparent = parentpath[:parentpath.rindex("_")] + "_gene"
    os.makedirs(outparent, exist_ok=True)

    print(abspath)
    print(parentpath)
    print(basename)
    print(outparent)

    total_fasta = []

    for seq_record in SeqIO.parse(abspath, "fasta"):
        total_fasta += GenerateFasta(collection, seq_record)
        print(seq_record.id + " handled.")

    SeqIO.write(total_fasta, outparent + "/" + basename, "fasta")

    CloseMongoDB(client)


if __name__ == '__main__':
    path = os.path.abspath(sys.argv[1])
    print(path)
    for file in os.listdir(path):
        abspath = os.path.join(path, file)
        GenerateFile(abspath)
        print(abspath + " wrote.")
